import datetime

import pandas as pd
import matplotlib.pyplot as plt
import os
import multiprocessing
import threading
import time


def plot_navs(frequency, begin_date, M=None, N=None):
    """绘制净值曲线"""
    if M is None and N is None:
        files = os.listdir(f'{frequency}/{begin_date}/nav')
    elif N is None:
        files = os.listdir(f'{frequency}/{begin_date}/nav')
        files = [file for file in files if int(file.split('_')[1]) in M]
    else:
        files = [f'nav_{int(M)}_{float(N)}.csv']
    for file in files:
        print(frequency, begin_date, file)
        nav_df = pd.read_csv(f'{frequency}/{begin_date}/nav/{file}')
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

        # 净值曲线
        ax1.plot(nav_df['date'], nav_df['nav0'], label=f'nav', color='purple')
        ax1.plot(nav_df['date'], nav_df['nav2'], label=f'nav_after_fee', color='blue')
        ax1.plot(nav_df['date'], nav_df['nav(hold_btc)'], label='Hold BTC', color='red')
        ax1.set_title('Strategy Performance')
        ax1.set_ylabel('NAV')
        ax1.set_xlabel('Date')
        ax1.legend()

        ax2.plot(nav_df['date'], nav_df['nav0'], label=f'nav', color='purple')
        ax2.plot(nav_df['date'], nav_df['nav2'], label=f'nav_after_fee', color='blue')
        ax2.plot(nav_df['date'], nav_df['nav(hold_btc)'], label='Hold BTC', color='red')
        ax2.set_yscale('log')  # 设置 y 轴为对数尺度
        ax2.set_title('Strategy Performance (Logarithmic Scale)')
        ax2.set_ylabel('NAV(log)')
        ax2.set_xlabel('Date')
        ax2.legend()

        plt.tight_layout()
        plt.savefig(f'{frequency}/{begin_date}/png/{file[:-4]}.png')
        plt.close(fig)  # 每次循环后关闭


def split_list_equal(lst, n=4):
    """将列表尽可能平均分成n组"""
    k, m = divmod(len(lst), n)
    return [lst[i*k+min(i, m):(i+1)*k+min(i+1, m)] for i in range(n)]


def print_fulfillment_of_schedule():
    total_files_num = 0
    for frequency in ["1d", "3h", "1h"]:
        for begin_date in ["20170101", "20200101", "20230101"]:
            files = os.listdir(f'{frequency}/{begin_date}/nav')
            total_files_num += len(files)
    while True:
        totoal_png_num = 0
        for frequency in ["1d", "3h", "1h"]:
            for begin_date in ["20170101", "20200101", "20230101"]:
                files = os.listdir(f'{frequency}/{begin_date}/png')
                totoal_png_num += len(files)
        print(f"{datetime.datetime.now()},进度:文件总数:{total_files_num},png总数:{totoal_png_num}")
        time.sleep(10)


if __name__ == '__main__':
    # 指定参数画净值图
    # plot_navs(frequency='1d', begin_date='20170101', M=7)
    # 画全部净值图
    threading.Thread()
    daemon_thread = threading.Thread(
        target=print_fulfillment_of_schedule,
        daemon=True  # 设置为守护线程
    )

    # 启动线程
    daemon_thread.start()
    params = []
    from backtest_boll_breakup_with_atr_v2 import BollBreakupWithAtrV2

    for frequency in ["1d", "3h", "1h"]:
        obj = BollBreakupWithAtrV2(frequency)
        m_list = obj.M_list
        grps = split_list_equal(m_list, 12)
        for grp in grps:
            for begin_date in ["20170101", "20200101", "20230101"]:
                params.append((frequency, begin_date, grp))

    with multiprocessing.Pool() as pool:
        pool.starmap(plot_navs, params)
